package com.rain.rainspringboot3.encryption.strategy;

import com.rain.rainspringboot3.encryption.exception.DecryptionException;
import com.rain.rainspringboot3.encryption.exception.EncryptionException;
import com.rain.rainspringboot3.encryption.service.EncryptionServiceImpl;
import jakarta.annotation.PostConstruct;
import org.apache.commons.lang3.StringUtils;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;

import javax.crypto.Cipher;
import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import javax.crypto.spec.GCMParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.nio.charset.StandardCharsets;
import java.security.SecureRandom;
import java.util.Base64;

import static com.rain.rainspringboot3.encryption.service.EncryptionServiceImpl.KEY_CACHE_MAP;

/**
 * AES加密策略实现
 * 使用AES-GCM算法提供加密和解密功能
 */
@Component
public class AESEncryptionStrategy implements EncryptionStrategy {

    private static final String ALGORITHM = "AES/GCM/NoPadding";
    private static final int TAG_LENGTH_BIT = 128;
    private static final int IV_LENGTH_BYTE = 12;
    private static final int KEY_LENGTH_BIT = 256;

    @Value("${encryption.strategies.aes.iv}")
    private String defaultIv;

    @Value("${encryption.strategies.aes.key}")
    private String key;

    @Override
    public String encrypt(String content, String key) {
        try {
            if (StringUtils.isBlank(key)) {
                key = this.key;
            }
            // 生成随机IV
            byte[] iv = new byte[IV_LENGTH_BYTE];
            SecureRandom random = new SecureRandom();
            random.nextBytes(iv);

            // 创建密钥规范
            SecretKeySpec keySpec = new SecretKeySpec(Base64.getDecoder().decode(key), "AES");

            // 初始化加密器
            Cipher cipher = Cipher.getInstance(ALGORITHM);
            GCMParameterSpec parameterSpec = new GCMParameterSpec(TAG_LENGTH_BIT, iv);
            cipher.init(Cipher.ENCRYPT_MODE, keySpec, parameterSpec);

            // 执行加密
            byte[] encrypted = cipher.doFinal(content.getBytes(StandardCharsets.UTF_8));

            // 组合IV和加密结果
            byte[] combined = new byte[iv.length + encrypted.length];
            System.arraycopy(iv, 0, combined, 0, iv.length);
            System.arraycopy(encrypted, 0, combined, iv.length, encrypted.length);

            return Base64.getEncoder().encodeToString(combined);
        } catch (Exception e) {
            throw new EncryptionException("AES encryption failed", e);
        }
    }

    @Override
    public String decrypt(String content, String key) {
        try {
            // 解码Base64
            byte[] combined = Base64.getDecoder().decode(content);

            // 分离IV和加密数据
            byte[] iv = new byte[IV_LENGTH_BYTE];
            byte[] encrypted = new byte[combined.length - IV_LENGTH_BYTE];
            System.arraycopy(combined, 0, iv, 0, iv.length);
            System.arraycopy(combined, iv.length, encrypted, 0, encrypted.length);

            // 创建密钥规范
            SecretKeySpec keySpec = new SecretKeySpec(Base64.getDecoder().decode(key), "AES");

            // 初始化解密器
            Cipher cipher = Cipher.getInstance(ALGORITHM);
            GCMParameterSpec parameterSpec = new GCMParameterSpec(TAG_LENGTH_BIT, iv);
            cipher.init(Cipher.DECRYPT_MODE, keySpec, parameterSpec);

            // 执行解密
            byte[] decrypted = cipher.doFinal(encrypted);
            return new String(decrypted, StandardCharsets.UTF_8);
        } catch (Exception e) {
            throw new DecryptionException("AES decryption failed", e);
        }
    }

    @Override
    public String generateKey() {
        try {
            KeyGenerator keyGenerator = KeyGenerator.getInstance("AES");
            keyGenerator.init(KEY_LENGTH_BIT);
            SecretKey secretKey = keyGenerator.generateKey();
            return Base64.getEncoder().encodeToString(secretKey.getEncoded());
        } catch (Exception e) {
            throw new EncryptionException("Failed to generate AES key", e);
        }
    }

    @PostConstruct
    public void saveKey(){
        KEY_CACHE_MAP.put("AES",key);
    }
}